{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [
        "Led8SgZ0YnGB",
        "pYZtXvtmes2I",
        "bSug2SNMq9Hg"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "![visitors](https://visitor-badge.glitch.me/badge?page_id=wsh.kohya_train_webui) \n",
        "[![Visitors](https://api.visitorbadge.io/api/combined?path=wsh.kohya_train_webui&countColor=%232ccce4&style=flat&labelStyle=none)](https://visitorbadge.io/status?path=wsh.kohya_train_webui)\n",
        "[![GitHub Repo stars](https://img.shields.io/github/stars/WSH032/kohya-config-webui?style=social)](https://github.com/WSH032/kohya-config-webui)\n",
        "\n",
        "| Notebook Name | Description | Link | Old-Version |\n",
        "| --- | --- | --- | --- |\n",
        "| [Colab_Lora_train](https://github.com/WSH032/lora-scripts/) | 基于[Akegarasu/lora-scripts](https://github.com/Akegarasu/lora-scripts)的定制化Colab notebook | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/WSH032/lora-scripts/blob/main/Colab_Lora_train.ipynb) | [![](https://img.shields.io/static/v1?message=Older%20Version&logo=googlecolab&labelColor=5c5c5c&color=e74c3c&label=%20&style=flat)](https://colab.research.google.com/drive/1_f0qJdM43BSssNJWtgjIlk9DkIzLPadx) | \n",
        "| [kohya_train_webui](https://github.com/WSH032/kohya-config-webui) `NEW` | 基于[WSH032/kohya-config-webui](https://github.com/WSH032/kohya-config-webui)的WebUI版Colab notebook | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/WSH032/kohya-config-webui/blob/main/kohya_train_webui.ipynb) |\n",
        "\n",
        "如果你觉得此项目有用，可以去 [![GitHub Repo stars](https://img.shields.io/github/stars/WSH032/kohya-config-webui?style=social)](https://github.com/WSH032/kohya-config-webui)</a>  点一颗小星星，非常感谢你⭐\n",
        "\n",
        "---\n",
        "\n",
        "- [📚notebook的操作手册](https://www.bilibili.com/read/cv23401664)\n",
        "\n",
        "- 参数：\n",
        "\n",
        " - [🥶冷门而有用的参数](https://www.bilibili.com/video/BV1mo4y1t7Zu/)\n",
        " - [🆕新版参数](https://www.bilibili.com/video/BV13s4y1377X/)\n",
        "\n",
        "---\n",
        "\n",
        "Based on the work of [kohya-ss](https://github.com/kohya-ss/sd-scripts) and [Linaqruf](https://github.com/Linaqruf/kohya-trainer)\n",
        "\n",
        "WebUI from [WSH032](https://github.com/WSH032/kohya-config-webui)\n"
      ],
      "metadata": {
        "id": "ll6PRAEKIfjT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 更新日志\n",
        "\n",
        "> 2023年5月24日：适配colab的torch==2.0.1\n",
        "> \n",
        "> 内容：适合torch==2.0.1的xformer以发布，不再强制安装torch==2.0.0，依赖安装时间恢复至2分40秒\n"
      ],
      "metadata": {
        "id": "TO3SVJYKmQhT"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# （一）环境配置"
      ],
      "metadata": {
        "id": "Led8SgZ0YnGB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##初始化常量与挂载谷歌硬盘（只要重启过colab就要再运行一次）\n",
        "\n",
        "#@markdown 是否挂载谷歌硬盘（推荐）\n",
        "use_google_drive = True #@param {type:\"boolean\"}\n",
        "\n",
        "import os\n",
        "import shutil\n",
        "import sys\n",
        "from google.colab import drive\n",
        "\n",
        "\n",
        "ROOT_DIR = os.getcwd()    #获取根目录\n",
        "\n",
        "SD_SCRIPTS_DIR = os.path.join( ROOT_DIR, \"sd-scripts\" )    #kohya库克隆路径\n",
        "WEBUI_DIR = os.path.join( ROOT_DIR, \"kohya-config-webui\" )   #webui库克隆路径\n",
        "\n",
        "#TRAIN_DATA_DIR = os.path.join( ROOT_DIR, \"Lora\", \"input\" )    #拷贝后训练材料路径\n",
        "#REG_DATA_DIR = os.path.join( ROOT_DIR, \"Lora\", \"reg\" )   #拷贝后正则化材料路径\n",
        "\n",
        "SD_MODEL_DIR = os.path.join( ROOT_DIR, \"Lora\", \"sd_model\" )    #SD模型下载地址\n",
        "VAE_MODEL_DIR = os.path.join( ROOT_DIR, \"Lora\", \"vae_model\" )    #VAE模型下载地址\n",
        "\n",
        "DEFAULT_COLAB_INPUT_DIR = os.path.normpath(\"/content/drive/MyDrive/Lora/input\")    #默认Colab训练集地址\n",
        "DEFAULT_COLAB_REG_DIR = os.path.normpath(\"/content/drive/MyDrive/Lora/reg\")    #默认Colab正则化地址\n",
        "DEFAULT_COLAB_OUPUT_DIR = os.path.normpath(\"/content/drive/MyDrive/Lora/output\")    #默认Colab模型输出地址\n",
        "DEFAULT_COLAB_WEBUI_SAVE_DIR = os.path.normpath(\"/content/drive/MyDrive/Lora/kohya_config_webui_save\")    #默认Colab保存webui参数文件地址\n",
        "\n",
        "ACCELERATE_CONFIG_PATH = os.path.join( ROOT_DIR, \"accelerate_config.yaml\" )   #accelerate库config文件写入地址\n",
        "\n",
        "\n",
        "#@title ##挂载谷歌硬盘\n",
        "\n",
        "if use_google_drive:\n",
        "    if not os.path.exists(\"/content/drive\"):\n",
        "        drive.mount(\"/content/drive\")\n",
        "\n",
        "!nvidia-smi\n",
        "\n",
        "#训练用环境变量\n",
        "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
        "os.environ[\"BITSANDBYTES_NOWELCOME\"] = \"1\"  \n",
        "os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\""
      ],
      "metadata": {
        "id": "lcFMoxnjwCDV",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##克隆github的库、安装依赖\n",
        "os.chdir( ROOT_DIR )\n",
        "!git clone https://github.com/kohya-ss/sd-scripts.git {SD_SCRIPTS_DIR}\n",
        "#@title 克隆我的库 \n",
        "!git clone https://github.com/WSH032/kohya-config-webui.git {WEBUI_DIR}\n",
        "\n",
        "#安装torch\n",
        "print(f\"torch安装中\")\n",
        "!pip -q install torch torchvision xformers triton\n",
        "print(f\"torch安装完成\")\n",
        "\n",
        "#安装kohya依赖\n",
        "print(f\"kohya依赖安装中\")\n",
        "os.chdir(SD_SCRIPTS_DIR)\n",
        "!pip -q install -r requirements.txt\n",
        "os.chdir(ROOT_DIR)\n",
        "print(f\"kohya依赖安装完成\")\n",
        "\n",
        "#安装lion优化器、Dadaption优化器、lycoris\n",
        "print(f\"lion优化器、Dadaption优化器、lycoris安装中\")\n",
        "!pip -q install --upgrade lion-pytorch dadaptation lycoris-lora\n",
        "print(f\"lion优化器、Dadaption优化器、lycoris安装完成\")\n",
        "\n",
        "#安装wandb\n",
        "print(f\"wandb安装中\")\n",
        "!pip -q install wandb\n",
        "print(f\"wandb安装中\")\n",
        "\n",
        "#安装webui依赖\n",
        "print(f\"webui依赖安装中\")\n",
        "os.chdir(WEBUI_DIR)\n",
        "!pip -q install -r requirements.txt\n",
        "os.chdir(ROOT_DIR)\n",
        "print(f\"webui依赖安装完成\")\n",
        "\n",
        "#安装功能性依赖\n",
        "!apt -q install aria2\n",
        "!pip -q install portpicker\n",
        "\n",
        "\n",
        "import torch\n",
        "print(\"当前torch版本\",torch.__version__)\n",
        "import torchvision\n",
        "print(\"当前torchvision版本\",torchvision.__version__)\n",
        "import triton\n",
        "print(\"当前triton版本\", triton.__version__)\n",
        "\n",
        "!python -V"
      ],
      "metadata": {
        "cellView": "form",
        "id": "4wYnLUrYY6Ut"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 下载模型, 可以同时选多个模型下载，到时候是在WebUI里选（原始代码来源于：[Linaqruf](https://github.com/Linaqruf/kohya-trainer)）\n",
        "installModels = []\n",
        "installv2Models = []\n",
        "\n",
        "#@markdown **预设底模**\n",
        "\n",
        "#@markdown SD1.x model\n",
        "modelName = \"Animefull-final-pruned.ckpt\"  # @param [\"\", \"Animefull-final-pruned.ckpt\", \"Anything-v3-1.safetensors\", \"AnyLoRA.safetensors\", \"AnimePastelDream.safetensors\", \"Chillout-mix.safetensors\", \"OpenJourney-v4.ckpt\", \"Stable-Diffusion-v1-5.safetensors\"]\n",
        "#@markdown SD2.x model `这些为SD2.x模型，训练时请开启v2选项`\n",
        "v2ModelName = \"\"  # @param [\"\", \"stable-diffusion-2-1-base.safetensors\", \"stable-diffusion-2-1-768v.safetensors\", \"plat-diffusion-v1-3-1.safetensors\", \"replicant-v1.safetensors\", \"illuminati-diffusion-v1-0.safetensors\", \"illuminati-diffusion-v1-1.safetensors\", \"waifu-diffusion-1-4-anime-e2.ckpt\", \"waifu-diffusion-1-5-e2.safetensors\", \"waifu-diffusion-1-5-e2-aesthetic.safetensors\"]\n",
        "\n",
        "#@markdown **自定义模型（不能超过5G）URL例如**`https://huggingface.co/a1079602570/animefull-final-pruned/resolve/main/novelailatest-pruned.ckpt`\n",
        "\n",
        "base_model_url = \"\" #@param {type:\"string\"}\n",
        "\n",
        "#@markdown **或者自定义模型（不能超过5G）路径例如**`/content/drive/MyDrive/Lora/model/your_model.ckpt`\n",
        "\n",
        "base_model_self_path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "\n",
        "def get_sd_model():\n",
        "    modelUrl = [\n",
        "        \"\",\n",
        "        \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/models/animefull-final-pruned.ckpt\",\n",
        "        \"https://huggingface.co/cag/anything-v3-1/resolve/main/anything-v3-1.safetensors\",\n",
        "        \"https://huggingface.co/Lykon/AnyLoRA/resolve/main/AnyLoRA_noVae_fp16.safetensors\",\n",
        "        \"https://huggingface.co/Lykon/AnimePastelDream/resolve/main/AnimePastelDream_Soft_noVae_fp16.safetensors\",\n",
        "        \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/chillout_mix-pruned.safetensors\",\n",
        "        \"https://huggingface.co/prompthero/openjourney-v4/resolve/main/openjourney-v4.ckpt\",\n",
        "        \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/stable_diffusion_1_5-pruned.safetensors\",\n",
        "    ]\n",
        "    modelList = [\n",
        "        \"\",\n",
        "        \"Animefull-final-pruned.ckpt\",\n",
        "        \"Anything-v3-1.safetensors\",\n",
        "        \"AnyLoRA.safetensors\",\n",
        "        \"AnimePastelDream.safetensors\",    \n",
        "        \"Chillout-mix.safetensors\",\n",
        "        \"OpenJourney-v4.ckpt\",\n",
        "        \"Stable-Diffusion-v1-5.safetensors\",\n",
        "    ]\n",
        "    v2ModelUrl = [\n",
        "        \"\",\n",
        "        \"https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.safetensors\",\n",
        "        \"https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors\",\n",
        "        \"https://huggingface.co/p1atdev/pd-archive/resolve/main/plat-v1-3-1.safetensors\",\n",
        "        \"https://huggingface.co/gsdf/Replicant-V1.0/resolve/main/Replicant-V1.0.safetensors\",\n",
        "        \"https://huggingface.co/IlluminatiAI/Illuminati_Diffusion_v1.0/resolve/main/illuminati_diffusion_v1.0.safetensors\",\n",
        "        \"https://huggingface.co/4eJIoBek/Illuminati-Diffusion-v1-1/resolve/main/illuminatiDiffusionV1_v11.safetensors\",\n",
        "        \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/wd-1-4-anime_e2.ckpt\",\n",
        "        \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-fp32.safetensors\",\n",
        "        \"https://huggingface.co/waifu-diffusion/wd-1-5-beta2/resolve/main/checkpoints/wd-1-5-beta2-aesthetic-fp32.safetensors\",\n",
        "    ]\n",
        "    v2ModelList = [\n",
        "        \"\",\n",
        "        \"stable-diffusion-2-1-base.safetensors\",\n",
        "        \"stable-diffusion-2-1-768v.safetensors\",\n",
        "        \"plat-diffusion-v1-3-1.safetensors\",\n",
        "        \"replicant-v1.safetensors\",\n",
        "        \"illuminati-diffusion-v1-0.safetensors\",\n",
        "        \"illuminati-diffusion-v1-1.safetensors\",\n",
        "        \"waifu-diffusion-1-4-anime-e2.ckpt\",\n",
        "        \"waifu-diffusion-1-5-e2.safetensors\",\n",
        "        \"waifu-diffusion-1-5-e2-aesthetic.safetensors\",\n",
        "    ]\n",
        "    if modelName:\n",
        "        installModels.append((modelName, modelUrl[modelList.index(modelName)]))\n",
        "    if v2ModelName:\n",
        "        installv2Models.append((v2ModelName, v2ModelUrl[v2ModelList.index(v2ModelName)]))\n",
        "\n",
        "\n",
        "    #下载模型\n",
        "    def install(checkpoint_name, url):\n",
        "        hf_token = \"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\"\n",
        "        user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "        print(checkpoint_name)\n",
        "        print(url)\n",
        "        !aria2c --console-log-level=error --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {SD_MODEL_DIR} -o {checkpoint_name} {url}\n",
        "    def install_checkpoint():\n",
        "        for model in installModels:\n",
        "            install(model[0], model[1])\n",
        "        for v2model in installv2Models:\n",
        "            install(v2model[0], v2model[1])\n",
        "\n",
        "    #下载预设模型\n",
        "    install_checkpoint()\n",
        "\n",
        "    #自定义链接不留空，则尝试下载\n",
        "    if base_model_url:\n",
        "        #!aria2c --content-disposition-default-utf8=true --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d {SD_MODEL_DIR} {base_model_url}\n",
        "        !wget {base_model_url} -P {SD_MODEL_DIR} -N\n",
        "\n",
        "    #自定义路径不留空，则尝试拷贝\n",
        "    if base_model_self_path:\n",
        "        try:\n",
        "            base_model_copy_path = os.path.join( SD_MODEL_DIR, os.path.basename(base_model_self_path) )\n",
        "            shutil.copyfile(base_model_self_path, base_model_copy_path)\n",
        "            print(f\"拷贝自定义底模成功, {base_model_self_path}被拷贝至{base_model_copy_path}\")\n",
        "        except Exception as e:\n",
        "            print(f\"拷贝自定义底模时发生错误， Error: {e}\")\n",
        "\n",
        "get_sd_model()\n",
        "\n",
        "\n",
        "#@markdown **(可选)选择一个Vae下载**`\"animevae.pt\", \"kl-f8-anime.ckpt\", \"vae-ft-mse-840000-ema-pruned.ckpt\"`\n",
        "\n",
        "vaeName = \"\"  # @param [\"\", \"anime.vae.pt\", \"waifudiffusion.vae.pt\", \"stablediffusion.vae.pt\"]\n",
        "\n",
        "def get_vae_model():\n",
        "\n",
        "    installVae = []\n",
        "\n",
        "    vaeUrl = [\n",
        "        \"\",\n",
        "        \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/vae/animevae.pt\",\n",
        "        \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime.ckpt\",\n",
        "        \"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt\",\n",
        "    ]\n",
        "    vaeList = [\"\", \"anime.vae.pt\", \"waifudiffusion.vae.pt\", \"stablediffusion.vae.pt\"]\n",
        "\n",
        "    installVae.append((vaeName, vaeUrl[vaeList.index(vaeName)]))\n",
        "\n",
        "    #开始下载\n",
        "    def install(vae_name, url):\n",
        "        hf_token = \"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\"\n",
        "        user_header = f'\"Authorization: Bearer {hf_token}\"'\n",
        "        print(vae_name)\n",
        "        print(url)\n",
        "        !aria2c --console-log-level=error --allow-overwrite --summary-interval=10 --header={user_header} -c -x 16 -k 1M -s 16 -d {VAE_MODEL_DIR} -o {vae_name} \"{url}\"\n",
        "\n",
        "    def install_vae():\n",
        "        if vaeName:\n",
        "            for vae in installVae:\n",
        "                install(vae[0], vae[1])\n",
        "        else:\n",
        "            pass\n",
        "    install_vae()\n",
        "\n",
        "get_vae_model()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "OYGUN309MuUJ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# （二）训练参数"
      ],
      "metadata": {
        "id": "pYZtXvtmes2I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##启动WebUI来设置参数\n",
        "\n",
        "#@markdown - 在谷歌硬盘的`/content/drive/MyDrive/Lora/kohya_config_webui_save`会生成一个`colab.toml`，在WebUI里读取它，会帮你完成默认参数设置。\n",
        "#@markdown  - 读取的时候会提示参数找不到，这是正常的\n",
        "#@markdown - 设置好参数后可以保存`（默认会保存到你的谷歌硬盘）`，以后读取你保存的配置文件就行\n",
        "#@markdown  - 保存toml配置文件时候不要用`colab.toml`这个名字，会被覆盖掉\n",
        "\n",
        "#@markdown - 在colab里要开`lowram`，不然很多模型载入不了，读取`colab.toml`的时候会自动帮你开启\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown 是否在colab里打开webui`不勾选就输出一个链接，点击后在另一个网页操作，反正我喜欢不勾选`\n",
        "in_colab = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown 是否使用gradio的远程分享及队列功能\n",
        "use_queue = False #@param {type:\"boolean\"}\n",
        "\n",
        "#生成一个colab默认toml文件\n",
        "def creat_save_toml(save_dir):\n",
        "    \"\"\"生成适用于Colab的webui参数保存文件colab.toml\"\"\"\n",
        "    import toml\n",
        "    #写入路径\n",
        "    other={\"write_files_dir\":SD_SCRIPTS_DIR}\n",
        "    #材料、模型、输出路径\n",
        "    param={\n",
        "        \"train_data_dir\":DEFAULT_COLAB_INPUT_DIR,\n",
        "        \"reg_data_dir\":DEFAULT_COLAB_REG_DIR,\n",
        "        \"base_model_dir\":SD_MODEL_DIR,\n",
        "        \"vae_model_dir\":VAE_MODEL_DIR,\n",
        "        \"output_dir\":DEFAULT_COLAB_OUPUT_DIR,\n",
        "        \"lowram\":True,\n",
        "    }\n",
        "\n",
        "    save_dict = {\"other\":other, \"param\":param}\n",
        "    #写入文件\n",
        "    save_name = \"colab.toml\"\n",
        "    save_path = os.path.join( save_dir, save_name )\n",
        "    os.makedirs(save_dir, exist_ok=True)\n",
        "    with open(save_path, \"w\", encoding=\"utf-8\") as f:\n",
        "        f.write( toml.dumps(save_dict) )\n",
        "\n",
        "creat_save_toml(DEFAULT_COLAB_WEBUI_SAVE_DIR)\n",
        "\n",
        "#导入并生成demo\n",
        "launch_param = [f\"--save_dir={DEFAULT_COLAB_WEBUI_SAVE_DIR}\",\n",
        "        f\"--save_name=kohya_config_webui_save.toml\",\n",
        "        f\"--read_dir={DEFAULT_COLAB_WEBUI_SAVE_DIR}\"\n",
        "]\n",
        "os.chdir( os.path.join(WEBUI_DIR, \"module\") )\n",
        "from kohya_config_webui import create_demo\n",
        "os.chdir(ROOT_DIR)\n",
        "demo = create_demo(launch_param)\n",
        "\n",
        "#找一个空闲端口\n",
        "import portpicker\n",
        "port = portpicker.pick_unused_port()\n",
        "#启动\n",
        "if not use_queue:\n",
        "    demo.launch(server_port=port, inbrowser=False, inline=False)\n",
        "    #暴露端口\n",
        "    from google.colab import output\n",
        "    output.serve_kernel_port_as_window(port)\n",
        "    #是否在Colab里打开\n",
        "    if in_colab:\n",
        "        output.serve_kernel_port_as_iframe(port)\n",
        "else:\n",
        "    demo.queue().launch(server_port=port, inline=in_colab)\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "AlRW5ufPM-0a"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title  ### 开始训练\n",
        "\n",
        "#@markdown 若正确运行，训练完成后，模型会自动保存至你在WebUI里设置的地址\n",
        "\n",
        "#@markdown 默认训练配置文件在 `/content/sd-scripts/config_file.toml`\n",
        "\n",
        "#@markdown 默认采样参数文件在 `/content/sd-scripts/sample_prompts.txt`\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "#@markdown 如果你想用自己的配置文件，或者采样文件，请填入下方 `填入意味着启用`\n",
        "\n",
        "config_file_self_path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "sample_prompts_self_path = \"\" #@param {type:\"string\"}\n",
        "\n",
        "os.chdir(ROOT_DIR)\n",
        "\n",
        "from accelerate.utils import write_basic_config\n",
        "if not os.path.exists(ACCELERATE_CONFIG_PATH):\n",
        "    write_basic_config(save_location=ACCELERATE_CONFIG_PATH)\n",
        "\n",
        "\n",
        "\n",
        "os.chdir(SD_SCRIPTS_DIR)\n",
        "\n",
        "#开始训练！\n",
        "!accelerate launch --config_file={ACCELERATE_CONFIG_PATH} --num_cpu_threads_per_process=8 train_network.py\\\n",
        "  --config_file={config_file_self_path if config_file_self_path else \"config_file.toml\"}\\\n",
        "  --sample_prompts={sample_prompts_self_path if sample_prompts_self_path else \"sample_prompts.txt\"}\n",
        "\n",
        "os.chdir(ROOT_DIR)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "NTRgMI7jR3DY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# （三）开发代码`别碰`"
      ],
      "metadata": {
        "id": "bSug2SNMq9Hg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title linaqfuf优化代码\n",
        "\n",
        "!sed -i \"s@cpu@cuda@\" /content/sd-scripts/library/model_util.py\n",
        "\n",
        "import zipfile\n",
        "def ubuntu_deps(url, name, dst):\n",
        "    !wget --show-progress {url}\n",
        "    with zipfile.ZipFile(name, \"r\") as deps:\n",
        "        deps.extractall(dst)\n",
        "    !dpkg -i {dst}/*\n",
        "    os.remove(name)\n",
        "    shutil.rmtree(dst)\n",
        "deps_dir = \"/conent/dep\"\n",
        "ubuntu_deps(\n",
        "    \"https://huggingface.co/Linaqruf/fast-repo/resolve/main/deb-libs.zip\",\n",
        "    \"deb-libs.zip\",\n",
        "    deps_dir,\n",
        ")\n",
        "\n",
        "!apt -y update\n",
        "!apt install libunwind8-dev\n",
        "\n",
        "os.environ[\"LD_PRELOAD\"] = \"libtcmalloc.so\"\n",
        "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
        "os.environ[\"BITSANDBYTES_NOWELCOME\"] = \"1\"  \n",
        "os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
        "\n",
        "cuda_path = \"/usr/local/cuda-11.8/targets/x86_64-linux/lib/\"\n",
        "ld_library_path = os.environ.get(\"LD_LIBRARY_PATH\", \"\")\n",
        "os.environ[\"LD_LIBRARY_PATH\"] = f\"{ld_library_path}:{cuda_path}\""
      ],
      "metadata": {
        "cellView": "form",
        "id": "ps6GgFwVaqXx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ##拷贝材料(支持重复训练时选择新的路径)\n",
        "\n",
        "#@markdown 训练集路径，正则化集路径(正则化留空则不拷贝)\n",
        "\n",
        "#@markdown `教程默认路径：`\n",
        "\n",
        "#@markdown `训练集：/content/drive/MyDrive/Lora/input/`\n",
        "\n",
        "#@markdown `正则化：/content/drive/MyDrive/Lora/reg/`\n",
        "\n",
        "train_data_dir_self = \"/content/drive/MyDrive/Lora/input/blue_archive\" #@param {type:'string'}\n",
        "reg_data_dir_self = \"\" #@param {type:'string'}\n",
        "\n",
        "\n",
        "def copy_data_and_reg(data_dir: str, reg_dir: str = \"\"):\n",
        "    \"\"\"\n",
        "    将材料拷贝至TRAIN_DATA_DIR和REG_DATA_DIR\n",
        "    拷贝前会删除之前材料\n",
        "    data_dir为训练集，必填； reg_dir，默认为空，不填则不拷贝\n",
        "    \"\"\"\n",
        "    #训练集路径为空直接退出\n",
        "    if not data_dir:\n",
        "        print(f\"训练集路径为空\")\n",
        "        return\n",
        "\n",
        "    #已经存在拷贝材料则删除\n",
        "    def rm_dir(dir):\n",
        "        if os.path.exists(dir):\n",
        "            shutil.rmtree(dir)\n",
        "    rm_dir(TRAIN_DATA_DIR)\n",
        "    rm_dir(REG_DATA_DIR)\n",
        "\n",
        "    #拷贝材料\n",
        "    def cp_dir(from_dir, to_dir, name):\n",
        "        print(f\"拷贝{name}中\")\n",
        "        try:\n",
        "            shutil.copytree(from_dir, to_dir, dirs_exist_ok=True)\n",
        "            print(f\"{name}拷贝成功, {from_dir}被拷贝至{to_dir}\")\n",
        "        except Exception as e:\n",
        "            print(f\"拷贝{name}时发生错误， Error: {e}\")\n",
        "\n",
        "    cp_dir(data_dir, TRAIN_DATA_DIR, \"训练集\")\n",
        "    if reg_dir:\n",
        "        cp_dir(reg_dir, REG_DATA_DIR, \"训练集\")\n",
        "    else:\n",
        "        print(f\"不拷贝正则化\")\n",
        "\n",
        "copy_data_and_reg(train_data_dir_self, reg_data_dir_self)\n"
      ],
      "metadata": {
        "id": "EzHADOPZMlN4",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}
